#!/usr/bin/env python

import csv
import pysam

from bam_counter.base_counter import BaseCounter
from bam_counter.paired_base_counter import PairedBaseCounter
from bam_counter.positions_iterator import PositionCountIterator, CharacterDelimitedPositionsReader
from bam_counter.ref_map import RefMapBaseCounter, load_ref_map

def paired_bams_to_tsv(args):
    dna_bam = pysam.Samfile(args.dna_bam_file_name, 'rb')
    rna_bam = pysam.Samfile(args.rna_bam_file_name, 'rb')
    
    dna_counter = BaseCounter(dna_bam, args.min_base_qual, args.min_map_qual)    
    rna_counter = BaseCounter(rna_bam, args.min_base_qual, args.min_map_qual)
    
    if args.map_file is not None:
        delimiter_char = get_delimiter_char(args.map_delimiter)
        
        ref_map = load_ref_map(args.map_file, delimiter_char)
        rna_counter = RefMapBaseCounter(rna_counter, ref_map)
    
    paired_counter = PairedBaseCounter(dna_counter, rna_counter)
    
    writer = open(args.out_file_name, 'w')
    
    if args.positions_file is None:
        write_all_positions(paired_counter, writer, args.min_depth)
    else:
        write_specified_positions(paired_counter, writer, args.min_depth, args.positions_file, args.positions_delimiter)
        
def write_all_positions(counter, writer, min_depth):
    for ref in sorted(counter.refs):
        for row in counter.iter_ref(ref):                       
            if row.depth >= min_depth:
                write_row(row, writer)

def write_specified_positions(counter, writer, min_depth, positions_file, delimiter):
    delimiter_char = get_delimiter_char(delimiter)
        
    positions_reader = CharacterDelimitedPositionsReader(positions_file, delimiter_char)
    
    positions_iter = PositionCountIterator(counter, positions_reader)
    
    for row in positions_iter:
        if row.depth >= min_depth:
            write_row(row, writer)

def write_row(row, writer):
    writer.write(str(row))
    writer.write("\n")
    
    print row

def get_delimiter_char(delimiter):
    if delimiter == 'tab':
        delimiter_char = '\t'
    elif delimiter == 'space':
        delimiter_char = ' '
    elif delimiter == 'comma':
        delimiter_char = ','
    
    return delimiter_char

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(prog='paired_dna_rna_seq_to_counts.py')
    
    parser.add_argument('dna_bam_file_name', help='''DNA BAM file.''')
    
    parser.add_argument('rna_bam_file_name', help='''RNA BAM file.''')
    
    parser.add_argument('out_file_name',
                         help='Name of tsv file to be created.')
    
    parser.add_argument('--min_depth', default=1, type=int,
                         help='''Minimum depth of coverage in both tumour and normal sample required to use a site in
                         analysis. Default is 1.''')
    
    parser.add_argument('--min_base_qual', default=10, type=int,
                         help='''Remove bases with base quality lower than this. Default is 10.''')
    
    parser.add_argument('--min_map_qual', default=10, type=int,
                         help='''Remove bases with mapping quality lower than this. Default is 10.''')
    
    parser.add_argument('--positions_file', default=None,
                        help='''Path to list of positions to create tsv on. Should be tab separated chrom pos.''')
    
    parser.add_argument('--positions_delimiter', choices=['tab', 'space', 'comma'], default='tab',
                        help='''Delimiter used to separate chrom and pos in positions file.''')
    
    parser.add_argument('--map_file', default=None,
                        help='''Path to list of positions to create tsv on. Should be tab separated chrom pos.''')
    
    parser.add_argument('--map_delimiter', choices=['tab', 'space', 'comma'], default='tab',
                    help='''Delimiter used to separate chrom and pos in positions file.''')
    
    parser.set_defaults(func=paired_bams_to_tsv)
    
    args = parser.parse_args()

    args.func(args)
